import json
import requests
from api_providers.baidu import BaiduTokenManager


class ProviderBaidu:
    def __init__(self, api_key, secret_key):
        self.api_key = api_key
        self.secret_key = secret_key
        self.access_token = BaiduTokenManager.get_token(self.api_key, self.secret_key)

    def translate(self, text, target_language):
        # 文心一言的翻译API示例，这里假设存在一个translate API endpoint
        translate_url = 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/completions/sqlcoder_7b'
        headers = {
            'Content-Type': 'application/json'
        }
        params = {
            'access_token': self.access_token
        }
        data = {
            'prompt': text,
        }
        response = requests.post(translate_url, headers=headers, params=params, data=json.dumps(data))
        if response:
            return response.json()
        return None

    def converse(self, conversation_id, message):
        chat_url = 'https://aip.baidubce.com/rpc/2.0/nlp/v1/ernie_bot'
        headers = {
            'Content-Type': 'application/json'
        }
        params = {
            'access_token': self.access_token
        }
        data = {
            'session_id': conversation_id,
            'query': message,
            'version': '3.0'  # 使用的文心一言版本
        }
        response = requests.post(chat_url, headers=headers, params=params, data=json.dumps(data))
        if response:
            return response.json().get('result', {}).get('response', '')
        return None
